package com.it.gmall.realtime.util;

import org.apache.kafka.common.TopicPartition;
import org.apache.spark.streaming.kafka010.OffsetRange;
import redis.clients.jedis.Jedis;

import java.util.HashMap;
import java.util.Map;
import java.util.Set;

/**
 * desc: 维护偏移量的工具类
 *
 * @author : code1997
 * @date :2021-02-2021/2/1 21:46
 */
public class OffsetManagerUtil {

    /**
     * 从redis中获取偏移量，type=hash
     * key:offset:topic:groupId
     * field:partition
     * value:偏移量
     *
     * @return ：偏移量
     */
    public static Map<TopicPartition, Long> getOffset(String topic, String groupId) {
        Jedis jedisClient = RedisUtil.getJedisClient();
        Map<TopicPartition, Long> offsets = new HashMap<>();
        String offsetKey = "offset:" + topic + ":" + groupId;
        Map<String, String> offsetMap = jedisClient.hgetAll(offsetKey);
        jedisClient.close();
        for (Map.Entry<String, String> entry : offsetMap.entrySet()) {
            String partition = entry.getKey();
            String offset = entry.getValue();
            System.out.println("正在读取分区：" + partition + "--->偏移量为：" + offset);
            offsets.put(new TopicPartition(topic, Integer.parseInt(partition)), Long.parseLong(offset));
        }

        return offsets;
    }

    public static void setOffset(String topic, String groupId, OffsetRange[] offsetRanges) {
        HashMap<String, String> offsetMap = new HashMap<>();
        for (OffsetRange offsetRange : offsetRanges) {
            int partition = offsetRange.partition();
            long fromOffset = offsetRange.fromOffset();
            long untilOffset = offsetRange.untilOffset();
            offsetMap.put(Integer.toString(partition), Long.toString(untilOffset));
            System.out.println("保存分区：" + partition + "从" + fromOffset + "---->" + untilOffset);
        }
        String offsetKey = "offset:" + topic + ":" + groupId;
        if (offsetMap != null && offsetMap.size() > 0) {
            Jedis jedisClient = RedisUtil.getJedisClient();
            jedisClient.hmset(offsetKey, offsetMap);
            jedisClient.close();
        }

    }

    public static void main(String[] args) {
        Map<TopicPartition, Long> offset = OffsetManagerUtil.getOffset("topic", "group001");
        for (Map.Entry<TopicPartition, Long> entry : offset.entrySet()) {
            System.out.println(entry.getKey() + "====" + entry.getValue());
        }
    }
}
